import torch
import numpy as np
from PIL import Image
from torchvision import transforms, models

model = models.vgg16(pretrained=True)
# 构建新的全连接层
model.classifier = torch.nn.Sequential(torch.nn.Linear(25088, 100),
                                       torch.nn.ReLU(),
                                       torch.nn.Dropout(p=0.5),
                                       torch.nn.Linear(100, 2))

model.load_state_dict(torch.load('cat_dog.pth'))

model.eval()

label = np.array(['cat', 'dog'])

# 数据预处理
transform = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor()
])


def predict(image_path):
    # 打开图片
    img = Image.open(image_path)
    # 数据处理，再增加一个维度
    img = transform(img).unsqueeze(0)
    # 预测得到结果
    outputs = model(img)
    # 获得最大值所在位置
    _, predicted = torch.max(outputs, 1)
    # 转化为类别名称
    print(label[predicted.item()])

predict('image/test/cat/cat.1000.jpg')

predict('image/test/cat/cat.1003.jpg')

predict('image/test/dog/dog.1006.jpg')

predict('image/test/dog/dog.1008.jpg')
